from flask import Blueprint, request, jsonify, current_app
from flask_jwt_extended import jwt_required, get_jwt_identity
from models import User, Message, DecryptionLog, db
from models.audit import AuditLog
from utils.auth import require_authentication, require_role, ROLE_LAW
from utils.encryption import EncryptionUtils
from utils.validation import validate_law_enforcement_request
from flask import current_app
import uuid
import datetime
from datetime import timezone
import hmac
import hashlib
from ipaddress import ip_address, ip_network
from sqlalchemy.sql.expression import or_
import secrets
import base64
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes
from utils.audit import log_audit_event

law_enforcement_bp = Blueprint('law_enforcement', __name__, url_prefix='/api/law')

def get_system_private_key():
    """安全获取系统私钥（支持HSM集成）"""
    try:
        import importlib
        key_manager = importlib.import_module('app.key_manager')
        get_private_key = key_manager.get_private_key
        return get_private_key("system_master")
    except Exception as e:
        current_app.logger.error(f"Failed to get system private key: {str(e)}")
        raise ValueError("System key not available") from e

def is_ip_whitelisted(ip, whitelist=None):
    """验证IP是否在白名单中（支持动态加载）"""
    if whitelist is None:
        whitelist = current_app.config.get('LAW_ENFORCEMENT_IP_WHITELIST', [])
    
    try:
        client = ip_address(ip)
        for cidr in whitelist:
            if client in ip_network(cidr):
                return True
        return False
    except ValueError:
        current_app.logger.warning(f"Invalid IP address: {ip}")
        return False

def generate_law_token(payload, secret=None):
    """生成执法HMAC令牌（带时间戳防重放）"""
    if secret is None:
        secret = current_app.config.get('LAW_ENFORCEMENT_SECRET')
    if not secret:
        raise ValueError("Law enforcement secret not configured")
    
    timestamp = str(int(datetime.datetime.now(timezone.utc).timestamp()))
    full_payload = f"{payload}|{timestamp}"
    token = hmac.new(
        secret.encode('utf-8'),
        full_payload.encode('utf-8'),
        hashlib.sha256
    ).hexdigest()
    return token, timestamp

def verify_law_token(token, payload, timestamp=None, max_age=300):
    """验证执法令牌（带时间戳和过期检查）"""
    secret = current_app.config.get('LAW_ENFORCEMENT_SECRET')
    if not secret:
        current_app.logger.error("Law enforcement secret not configured")
        return False
    
    try:
        if timestamp:
            current_time = int(datetime.datetime.now(timezone.utc).timestamp())
            if current_time - int(timestamp) > max_age:
                current_app.logger.warning(f"Law token expired: {timestamp}")
                return False
        
        expected_token, _ = generate_law_token(payload, secret)
        return hmac.compare_digest(token, expected_token)
    except Exception as e:
        current_app.logger.error(f"Law token verification error: {str(e)}")
        return False

def decrypt_master_key(user, system_key=None):
    """安全解密用户主密钥（支持密钥版本管理）"""
    if system_key is None:
        system_key = get_system_private_key()
    
    try:
        # 假设用户主密钥使用系统公钥加密存储
        encrypted_master_key = user.encrypted_master_key
        if isinstance(encrypted_master_key, str):
            encrypted_master_key = base64.b64decode(encrypted_master_key)
        
        # 支持多种padding方式（自动检测版本）
        try:
            master_key = system_key.decrypt(
                encrypted_master_key,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
        except:
            master_key = system_key.decrypt(
                encrypted_master_key,
                padding.PKCS1v15()
            )
        return master_key
    except Exception as e:
        logger.error(f"Master key decryption failed: {str(e)}")
        raise

@law_enforcement_bp.route('/decrypt', methods=['POST'])
@jwt_required()
@require_authentication
@require_role(ROLE_LAW)
def decrypt_message(current_user):
    """执法部门解密消息API（增强版）"""
    try:
        client_ip = request.remote_addr
        whitelist = current_app.config.get('LAW_ENFORCEMENT_IP_WHITELIST', [])
        if not is_ip_whitelisted(client_ip, whitelist):
            current_app.logger.warning(f"Unauthorized IP access: {client_ip}")
            log_audit_event(
                event_type='LAW_UNAUTHORIZED_IP',
                user_id=current_user.id,
                ip_address=client_ip,
                details=f"IP {client_ip} not in whitelist"
            )
            return jsonify({"error": "IP not authorized"}), 403
        
        data = request.get_json()
        validation = validate_law_enforcement_request(data)
        if validation.get("status") == "error":
            return jsonify(validation), 400
        
        # 增强版HMAC令牌验证（带时间戳防重放）
        message_id = data.get("message_id")
        reason = data.get("reason", "")
        authorized_by = data.get("authorized_by")
        token = data.get("authorization_token")
        timestamp = data.get("token_timestamp")
        
        request_payload = f"{message_id}|{reason}|{authorized_by}"
        if not verify_law_token(token, request_payload, timestamp):
            logger.warning(f"Invalid law token from {client_ip}")
            return jsonify({"error": "Invalid authorization token"}), 401
        
        # 数据存在性验证
        message_uuid = uuid.UUID(message_id)
        message = Message.query.get(message_uuid)
        if not message:
            return jsonify({"error": "Message not found"}), 404
        
        # 权限验证（执法部门只能解密指定范围内的消息）
        if not current_app.config.get("LAW_ENFORCEMENT_ENABLED", False):
            return jsonify({"error": "Law enforcement functionality disabled"}), 403
        
        # 获取系统密钥并解密用户主密钥
        system_key = get_system_private_key()
        target_user = User.query.get(message.sender_id) or User.query.get(message.receiver_id)
        if not target_user:
            return jsonify({"error": "Message owner not found"}), 404
        
        master_key = decrypt_master_key(target_user, system_key)
        
        # 解密消息内容
        try:
            decrypted_content = EncryptionUtils.aes_decrypt(
                master_key,
                message.encrypted_content,
                message.nonce,
                message.tag
            )
        except Exception as e:
            current_app.logger.error(f"Message decryption failed: {str(e)}")
            return jsonify({"error": "Failed to decrypt message content"}), 500
        
        # 记录脱敏后的解密日志
        decrypted_text = decrypted_content.decode('utf-8', errors='replace')
        masked_content = decrypted_text[:100] + ('...' if len(decrypted_text) > 100 else '')
        
        decryption_log = DecryptionLog(
            message_id=message_uuid,
            authorized_by=authorized_by,
            operator_id=current_user.id,
            operator_ip=client_ip,
            timestamp=datetime.datetime.now(timezone.utc),
            reason=reason,
            decrypted_content=base64.b64encode(decrypted_content).decode('utf-8'),
            request_payload=jsonify(data).data.decode('utf-8')
        )
        db.session.add(decryption_log)
        db.session.commit()
        
        # 多因素认证检查
        mfa_code = data.get('mfa_code')
        if not verify_law_enforcement_mfa(current_user.id, mfa_code):
            log_audit_event(
                event_type='LAW_MFA_FAILED',
                user_id=current_user.id,
                ip_address=client_ip,
                details=f"MFA verification failed for {current_user.username}"
            )
            return jsonify({"error": "Multi-factor authentication failed"}), 401

        # 记录详细审计日志
        log_audit_event(
            event_type='LAW_DECRYPT_SUCCESS',
            user_id=current_user.id,
            ip_address=client_ip,
            details=f"Decrypted message {message_id} for {authorized_by}"
        )

        return jsonify({
            "status": "success",
            "data": {
                "message_id": str(message_uuid),
                "sender_id": str(message.sender_id),
                "receiver_id": str(message.receiver_id),
                "timestamp": message.timestamp.isoformat(),
                "decrypted_content": masked_content,
                "decryption_log_id": str(decryption_log.id)
            }
        }), 200

    except ValueError:
        log_audit_event(
            event_type='LAW_INVALID_UUID',
            user_id=current_user.id,
            ip_address=client_ip,
            details=f"Invalid message ID format: {message_id}"
        )
        return jsonify({"error": "Invalid message ID format"}), 400
    except Exception as e:
        logger.error(f"Law decryption API error: {str(e)}", exc_info=True)
        log_audit_event(
            event_type='LAW_DECRYPT_ERROR',
            user_id=current_user.id,
            ip_address=client_ip,
            details=f"Decryption failed: {str(e)}"
        )
        return jsonify({
            "status": "error",
            "message": "Internal server error"
        }), 500

@law_enforcement_bp.route('/logs', methods=['GET'])
@jwt_required()
@require_authentication
@require_role(ROLE_LAW)
def get_decryption_logs(current_user):
    """获取解密日志（带分页和过滤）"""
    try:
        page = request.args.get('page', 1, type=int)
        per_page = request.args.get('per_page', 20, type=int)
        start_date = request.args.get('start_date')
        end_date = request.args.get('end_date')
        user_id = request.args.get('user_id')
        
        query = DecryptionLog.query.order_by(DecryptionLog.timestamp.desc())
        
        if start_date:
            query = query.filter(DecryptionLog.timestamp >= start_date)
        if end_date:
            query = query.filter(DecryptionLog.timestamp <= end_date)
        if user_id:
            try:
                user_uuid = uuid.UUID(user_id)
                message_query = Message.query.filter(
                    or_(Message.sender_id == user_uuid, Message.receiver_id == user_uuid)
                ).subquery()
                query = query.join(
                    message_query, 
                    DecryptionLog.message_id == message_query.c.id
                )
            except ValueError:
                pass
        
        logs = query.paginate(page=page, per_page=per_page, error_out=False)
        
        log_list = []
        for log in logs.items:
            log_data = {
                "id": str(log.id),
                "message_id": str(log.message_id),
                "authorized_by": log.authorized_by,
                "operator_id": str(log.operator_id),
                "operator_ip": log.operator_ip,
                "timestamp": log.timestamp.isoformat(),
                "reason": log.reason,
                "content_length": len(base64.b64decode(log.decrypted_content)) if log.decrypted_content else 0
            }
            if log.message:
                log_data["message_details"] = {
                    "sender_id": str(log.message.sender_id),
                    "receiver_id": str(log.message.receiver_id),
                    "timestamp": log.message.timestamp.isoformat()
                }
            log_list.append(log_data)
        
        return jsonify({
            "status": "success",
            "total": logs.total,
            "pages": logs.pages,
            "current_page": logs.page,
            "logs": log_list
        }), 200
        
    except Exception as e:
        logger.error(f"Get decryption logs error: {str(e)}", exc_info=True)
        return jsonify({
            "status": "error",
            "message": "Failed to get logs"
        }), 500

@law_enforcement_bp.route('/user/<string:user_id>/activity', methods=['GET'])
@jwt_required()
@require_authentication
@require_role(ROLE_LAW)
def get_user_activity(current_user, user_id):
    """获取用户活动记录（增强版）"""
    try:
        client_ip = request.remote_addr
        whitelist = current_app.config.get('LAW_ENFORCEMENT_IP_WHITELIST', [])
        if not is_ip_whitelisted(client_ip, whitelist):
            logger.warning(f"Unauthorized IP: {client_ip}")
            return jsonify({"error": "IP not authorized"}), 403
        
        user_uuid = uuid.UUID(user_id)
        user = User.query.get(user_uuid)
        if not user:
            return jsonify({"error": "User not found"}), 404
        
        # 消息统计（使用子查询优化性能）
        message_subq = db.session.query(
            Message.id,
            Message.sender_id,
            Message.receiver_id,
            Message.timestamp
        ).filter(
            or_(Message.sender_id == user_uuid, Message.receiver_id == user_uuid)
        ).subquery()
        
        total_messages = db.session.query(db.func.count(message_subq.c.id)).scalar()
        
        # 最近联系人（使用窗口函数优化）
        from sqlalchemy import func, desc, row_number
        from sqlalchemy.orm import aliased
        
        contact_subq = db.session.query(
            User.id,
            User.username,
            User.online,
            func.count(Message.id).label('message_count'),
            func.max(Message.timestamp).label('last_message'),
            row_number().over(
                partition_by=Message.id,
                order_by=func.max(Message.timestamp).desc()
            ).label('rn')
        ).join(
            Message,
            or_(User.id == Message.sender_id, User.id == Message.receiver_id)
        ).filter(
            or_(
                Message.sender_id == user_uuid,
                Message.receiver_id == user_uuid
            ),
            User.id != user_uuid
        ).group_by(User.id).subquery()
        
        recent_contacts = db.session.query(contact_subq).filter(contact_subq.c.rn <= 10).all()
        
        # 登录活动（使用日期桶优化）
        login_stats = db.session.query(
            func.date_trunc('day', AuthLog.timestamp).label('date'),
            func.count().label('login_count')
        ).filter(
            AuthLog.user_id == user_uuid,
            AuthLog.event_type == 'LOGIN_SUCCESS'
        ).group_by(func.date_trunc('day', AuthLog.timestamp)).order_by(
            func.date_trunc('day', AuthLog.timestamp).desc()
        ).limit(30).all()
        
        # 格式化响应数据
        return jsonify({
            "status": "success",
            "user": {
                "id": str(user_uuid),
                "username": user.username,
                "created_at": user.created_at.isoformat(),
                "last_login": user.last_login.isoformat() if user.last_login else None
            },
            "message_stats": {
                "total_messages": total_messages,
                "recent_contacts": [
                    {
                        "user_id": str(contact.id),
                        "username": contact.username,
                        "online": contact.online,
                        "message_count": contact.message_count,
                        "last_message": contact.last_message.isoformat() if contact.last_message else None
                    }
                    for contact in recent_contacts
                ]
            },
            "login_activity": [
                {
                    "date": log.date.strftime('%Y-%m-%d'),
                    "login_count": log.login_count
                }
                for log in login_stats
            ]
        }), 200
        
    except ValueError:
        return jsonify({"error": "Invalid user ID format"}), 400
    except Exception as e:
        logger.error(f"Get user activity error: {str(e)}", exc_info=True)
        return jsonify({
            "status": "error",
            "message": "Failed to get user activity"
        }), 500